
#include "CryptOChain.h"
#include "AESEncrypt.h"
#include "AESDecrypt.h"
#include "RSAEncrypt.h"
#include "RSADecrypt.h"
#include "RSASigning.h"
#include "RSAChecking.h"
#include "DSASigning.h"
#include "DSAChecking.h"
#include "ECDSASigning.h"
#include "ECDSAChecking.h"
#include "GZCompress.h"
#include "GZDecompress.h"

/** Constructor for the crypto chain. This is the constructor without
 * bucket brigades, so this can only manipulate blocks and not
 * streams.
 * 
 * @param imp The memory pool to use for allocation. */
CryptOChain::CryptOChain(apr_pool_t *imp)
{
   input = NULL;
   output = NULL;
   apr_status_t status;
   status = apr_pool_create(&mp, imp);
   CheckAPRError(status);
}

/** Constructor for the crypto chain. This is the constructor with
 * bucket brigades, so this can manipulate both blocks and streams.
 * 
 * @param iinput The bukcet brigade serving as input to the stream
 * process.
 * 
 * @param ioutput The bukcet brigade serving as output from the stream
 * process.
 * 
 * @param imp The memory pool to use for allocation. */
CryptOChain::CryptOChain(APRBucketBrigade &iinput, APRBucketBrigade &ioutput, apr_pool_t *imp)
{
   input = &iinput;
   output = &ioutput;
   apr_status_t status;
   status = apr_pool_create(&mp, imp);
   CheckAPRError(status);
   APRBucketBrigade *first = new APRBucketBrigade(mp);
   brigades.push_back(first);
}

/** Destructor for the crypto chain. */
CryptOChain::~CryptOChain()
{
   // Free the bucket brigades
   for (int i = 0; i < brigades.size(); i++)
      delete brigades.at(i);
   // Free the cryptopipes
   for (int i = 0; i < cryptopipes.size(); i++)
      delete cryptopipes.at(i);
   apr_pool_destroy(mp);
}

/** Adds an encryption step to the pipe or block process currently
 * represented by the class.
 * 
 * @param enc The encryption key with which to construct the new step
 * of the pipe or block process. */
void CryptOChain::addEncryptionPipe(EncryptionKey *enc)
{
   if (brigades.size() == 0)
   {
      CryptoPipeBlock *newcpb = NULL;
      switch (enc->getKeyType())
      {
         case AbstractKey::KEY_TYPE_PUBLIC_RSA:
            newcpb = (CryptoPipeBlock *)new RSAEncrypt(*((RSAPublicKey *)enc)); break;
         case AbstractKey::KEY_TYPE_SYMM_AES:
            newcpb = (CryptoPipeBlock *)new AESEncrypt(*((AESSymmetricKey *)enc)); break;
         default:
            errprint("unrecognised encryption key type"); break;
      }
      if (newcpb) cryptopipes.push_back(newcpb);
   } else
   {
      APRBucketBrigade *newbb = new APRBucketBrigade(mp);
      CryptoPipeBlock *newcpb = NULL;
      switch (enc->getKeyType())
      {
         case AbstractKey::KEY_TYPE_PUBLIC_RSA:
            newcpb = (CryptoPipeBlock *)new
               RSAEncrypt(*(brigades.at(brigades.size() - 1)), *newbb, *((RSAPublicKey *)enc)); break;
         case AbstractKey::KEY_TYPE_SYMM_AES:
            newcpb = (CryptoPipeBlock *)new
               AESEncrypt(*(brigades.at(brigades.size() - 1)), *newbb, *((AESSymmetricKey *)enc)); break;
         default:
            errprint("unrecognised encryption key type"); break;
      }
      if (newcpb)
      {
         brigades.push_back(newbb);
         cryptopipes.push_back(newcpb);
      } else delete newbb;
   }
}

/** Adds a decryption step to the pipe or block process currently
 * represented by the class.
 * 
 * @param dec The decryption key with which to construct the new step
 * of the pipe or block process. */
void CryptOChain::addDecryptionPipe(DecryptionKey *dec)
{
   if (brigades.size() == 0)
   {
      CryptoPipeBlock *newcpb = NULL;
      switch (dec->getKeyType())
      {
         case AbstractKey::KEY_TYPE_PRIVATE_RSA:
            newcpb = (CryptoPipeBlock *)new RSADecrypt(*((RSAPrivateKey *)dec)); break;
         case AbstractKey::KEY_TYPE_SYMM_AES:
            newcpb = (CryptoPipeBlock *)new AESDecrypt(*((AESSymmetricKey *)dec)); break;
         default:
            errprint("unrecognised decryption key type"); break;
      }
      if (newcpb) cryptopipes.push_back(newcpb);
   } else
   {
      APRBucketBrigade *newbb = new APRBucketBrigade(mp);
      CryptoPipeBlock *newcpb = NULL;
      switch (dec->getKeyType())
      {
         case AbstractKey::KEY_TYPE_PRIVATE_RSA:
            newcpb = (CryptoPipeBlock *)new
               RSADecrypt(*(brigades.at(brigades.size() - 1)), *newbb, *((RSAPrivateKey *)dec)); break;
         case AbstractKey::KEY_TYPE_SYMM_AES:
            newcpb = (CryptoPipeBlock *)new
               AESDecrypt(*(brigades.at(brigades.size() - 1)), *newbb, *((AESSymmetricKey *)dec)); break;
         default:
            errprint("unrecognised decryption key type"); break;
      }
      if (newcpb)
      {
         brigades.push_back(newbb);
         cryptopipes.push_back(newcpb);
      } else delete newbb;
   }
}

/** Adds a signing step to the pipe or block process currently
 * represented by the class.
 * 
 * @param sgn The signing key with which to construct the new step of
 * the pipe or block process. */
void CryptOChain::addSigningPipe(SigningKey *sgn)
{
   if (brigades.size() == 0)
   {
      CryptoPipeBlock *newcpb = NULL;
      switch (sgn->getKeyType())
      {
         case AbstractKey::KEY_TYPE_PRIVATE_RSA:
            newcpb = (CryptoPipeBlock *)new RSASigning(*((RSAPrivateKey *)sgn)); break;
         case AbstractKey::KEY_TYPE_PRIVATE_DSA:
            newcpb = (CryptoPipeBlock *)new DSASigning(*((DSAPrivateKey *)sgn)); break;
         case AbstractKey::KEY_TYPE_PRIVATE_ECDSA:
            newcpb = (CryptoPipeBlock *)new ECDSASigning(*((ECDSAPrivateKey *)sgn)); break;
         default:
            errprint("unrecognised signing key type"); break;
      }
      if (newcpb) cryptopipes.push_back(newcpb);
   } else
   {
      APRBucketBrigade *newbb = new APRBucketBrigade(mp);
      CryptoPipeBlock *newcpb = NULL;
      switch (sgn->getKeyType())
      {
         case AbstractKey::KEY_TYPE_PRIVATE_RSA:
            newcpb = (CryptoPipeBlock *)new
               RSASigning(*(brigades.at(brigades.size() - 1)), *newbb, *((RSAPrivateKey *)sgn)); break;
         case AbstractKey::KEY_TYPE_PRIVATE_DSA:
            newcpb = (CryptoPipeBlock *)new
               DSASigning(*(brigades.at(brigades.size() - 1)), *newbb, *((DSAPrivateKey *)sgn)); break;
         case AbstractKey::KEY_TYPE_PRIVATE_ECDSA:
            newcpb = (CryptoPipeBlock *)new
               ECDSASigning(*(brigades.at(brigades.size() - 1)), *newbb, *((ECDSAPrivateKey *)sgn)); break;
         default:
            errprint("unrecognised signing key type"); break;
      }
      if (newcpb)
      {
         brigades.push_back(newbb);
         cryptopipes.push_back(newcpb);
      } else delete newbb;
   }
}

/** Adds a checking step to the pipe or block process currently
 * represented by the class.
 * 
 * @param chk The checking key with which to construct the new step of
 * the pipe or block process. */
void CryptOChain::addCheckingPipe(CheckingKey *chk)
{
   if (brigades.size() == 0)
   {
      CryptoPipeBlock *newcpb = NULL;
      switch (chk->getKeyType())
      {
         case AbstractKey::KEY_TYPE_PUBLIC_RSA:
            newcpb = (CryptoPipeBlock *)new RSAChecking(*((RSAPublicKey *)chk)); break;
         case AbstractKey::KEY_TYPE_PUBLIC_DSA:
            newcpb = (CryptoPipeBlock *)new DSAChecking(*((DSAPublicKey *)chk)); break;
         case AbstractKey::KEY_TYPE_PUBLIC_ECDSA:
            newcpb = (CryptoPipeBlock *)new ECDSAChecking(*((ECDSAPublicKey *)chk)); break;
         default:
            errprint("unrecognised checking key type"); break;
      }
      if (newcpb) cryptopipes.push_back(newcpb);
   } else
   {
      APRBucketBrigade *newbb = new APRBucketBrigade(mp);
      CryptoPipeBlock *newcpb = NULL;
      switch (chk->getKeyType())
      {
         case AbstractKey::KEY_TYPE_PUBLIC_RSA:
            newcpb = (CryptoPipeBlock *)new
               RSAChecking(*(brigades.at(brigades.size() - 1)), *newbb, *((RSAPublicKey *)chk)); break;
         case AbstractKey::KEY_TYPE_PUBLIC_DSA:
            newcpb = (CryptoPipeBlock *)new
               DSAChecking(*(brigades.at(brigades.size() - 1)), *newbb, *((DSAPublicKey *)chk)); break;
         case AbstractKey::KEY_TYPE_PUBLIC_ECDSA:
            newcpb = (CryptoPipeBlock *)new
               ECDSAChecking(*(brigades.at(brigades.size() - 1)), *newbb, *((ECDSAPublicKey *)chk)); break;
         default:
            errprint("unrecognised checking key type"); break;
      }
      if (newcpb)
      {
         brigades.push_back(newbb);
         cryptopipes.push_back(newcpb);
      } else delete newbb;
   }
}

/** Adds a compression step to the pipe or block process currently
 * represented by the class. */
void CryptOChain::addCompressionPipe()
{
   if (brigades.size() == 0)
   {
      CryptoPipeBlock *newcpb = (CryptoPipeBlock *)new GZCompress();
      if (newcpb) cryptopipes.push_back(newcpb);
   } else
   {
      APRBucketBrigade *newbb = new APRBucketBrigade(mp);
      CryptoPipeBlock *newcpb = (CryptoPipeBlock *)new
         GZCompress(*(brigades.at(brigades.size() - 1)), *newbb);
      if (newcpb)
      {
         brigades.push_back(newbb);
         cryptopipes.push_back(newcpb);
      } else delete newbb;
   }
}

/** Adds a decompression step to the pipe or block process currently
 * represented by the class. */
void CryptOChain::addDecompressionPipe()
{
   if (brigades.size() == 0)
   {
      CryptoPipeBlock *newcpb = (CryptoPipeBlock *)new GZDecompress();
      if (newcpb) cryptopipes.push_back(newcpb);
   } else
   {
      APRBucketBrigade *newbb = new APRBucketBrigade(mp);
      CryptoPipeBlock *newcpb = (CryptoPipeBlock *)new
         GZDecompress(*(brigades.at(brigades.size() - 1)), *newbb);
      if (newcpb)
      {
         brigades.push_back(newbb);
         cryptopipes.push_back(newcpb);
      } else delete newbb;
   }
}

/** Generates a standard chain from encryption and signing.
 * 
 * @param crypto An encryption key, which, if specified is used for an
 * encryption step.
 * 
 * @param signing A signing key, which, if specified is used for a
 * signing step.
 * 
 * @param docompress Whether a compression step needs adding. */
void CryptOChain::buildEncryptionChain(EncryptionKey *crypto, SigningKey *signing, int docompress)
{
   if (docompress) addCompressionPipe();
   if (signing) addSigningPipe(signing);
   if (crypto) addEncryptionPipe(crypto);
}

/** Generates a standard chain from decryption and checking.
 * 
 * @param crypto An decryption key, which, if specified is used for a
 * decryption step.
 * 
 * @param signing A checking key, which, if specified is used for a
 * signature checking step.
 * 
 * @param dodecompress Whether a decompression step needs adding. */
void CryptOChain::buildDecryptionChain(DecryptionKey *crypto, CheckingKey *signing, int dodecompress)
{
   if (crypto) addDecryptionPipe(crypto);
   if (signing) addCheckingPipe(signing);
   if (dodecompress) addDecompressionPipe();
}

/** Process a block of data.
 * 
 * @param inputblock The buffer of the block of data to be processed.
 * 
 * @param inputsize The size of the block of data in the buffer.
 * 
 * @param outputsize The size of the output returned.
 * 
 * @return Returns the buffer of output. */
char *CryptOChain::processDataBlock(const char *inputblock, apr_size_t inputsize, apr_size_t *outputsize)
{
   const char *this_input_block = inputblock;
   apr_size_t this_input_size = inputsize;
   char *this_output_block = NULL;
   apr_size_t this_output_size = 0;
   for (int i = 0; i < cryptopipes.size(); i++)
   {
      this_output_block = cryptopipes.at(i)->processDataBlock(this_input_block, this_input_size, &this_output_size);
      // Only free it if it is not the original block (but must
      // const_cast!)
      if (this_input_block != inputblock)
         free(const_cast<char *>(this_input_block));
      this_input_block = this_output_block;
      this_input_size = this_output_size;
      if (this_output_block == NULL)
         return NULL;
   }
   trusted = 1;
   for (int i = 0; i < cryptopipes.size(); i++)
      trusted &= cryptopipes.at(i)->processWasSecure();
   if (outputsize)
      *outputsize = this_output_size;
   return this_output_block;
}

/** Begin processing the stream through the pipe.
 * 
 * @return Zero on success, non-zero on failure. */
int CryptOChain::processPipeBegin()
{
   for (int i = 0; i < cryptopipes.size(); i++)
   {
      cryptopipes.at(i)->processPipeBegin();
   }
   return 0;
}

/** Begin processing the stream through the pipe.
 * 
 * @param inputblocksize The block size to take from the input stream
 * and then place it into the chain.
 * 
 * @param outputblocksize The block size to take from the last link's
 * stream and put into the output stream.
 * 
 * @return Zero on completion, non-zero if the process is
 * incomplete. */
int CryptOChain::processPipeStep(int64_t inputblocksize, int64_t outputblocksize)
{
   char *data_block = (char *)malloc(sizeof(char)*inputblocksize);
   apr_size_t realdatasize = input->readFromStream(data_block, inputblocksize);
   apr_size_t realwrittensize = brigades.at(0)->writeToStream(data_block, realdatasize);
   int64_t total_in = realdatasize;
   int64_t total_out = 0;
   input->removeFromStream(realwrittensize);
   for (int i = 0; i < cryptopipes.size(); i++)
   {
      int input_change = 0;
      int output_change = 0;
      cryptopipes.at(i)->processPipeStep(&input_change, &output_change);
      total_in += input_change;
      total_out += output_change;
   }
   realdatasize = brigades.at(brigades.size() - 1)->readFromStream(data_block, outputblocksize);
   realwrittensize = output->writeToStream(data_block, realdatasize);
   brigades.at(brigades.size() - 1)->removeFromStream(realwrittensize);
   total_out += realwrittensize;
//   printf("in: %d out: %d\n", total_in, total_out);
   free(data_block);
   return ((total_in + total_out) > 0);
}

/** End processing the stream through the pipe.
 * 
 * @return Zero on success, non-zero on failure. */
int CryptOChain::processPipeEnd()
{
   while (processPipeStep(0x10000, 0x10000));
   for (int i = 0; i < cryptopipes.size(); i++)
   {
      cryptopipes.at(i)->processPipeEnd();
   }
   apr_size_t blocksize = brigades.at(brigades.size() - 1)->realLength();
   char *data_block = (char *)malloc(sizeof(char)*blocksize);
   apr_size_t realdatasize = brigades.at(brigades.size() - 1)->readFromStream(data_block, blocksize);
   apr_size_t realwrittensize = output->writeToStream(data_block, realdatasize);
   brigades.at(brigades.size() - 1)->removeFromStream(realwrittensize);
   trusted = 1;
   for (int i = 0; i < cryptopipes.size(); i++)
      trusted &= cryptopipes.at(i)->processWasSecure();
   //for (int i = 0; i < brigades.size(); i++)
   //   printf("brigade %d - %d bytes\n", i, brigades.at(i)->realLength());
   return 0;
}
